import numpy as np
from pylearn2.datasets.dense_design_matrix import DenseDesignMatrix
from pylearn2.utils import serial

class PylearnDataset(DenseDesignMatrix):
    
    def __init__(self, X, Y, preprocessor=None):
        one_hot = np.zeros(( len(Y), 2 ))
        for i in xrange(len(Y)):
            one_hot[i, Y[i]] = 1
            
        super(PylearnDataset, self).__init__(X=X, y=one_hot, preprocessor=preprocessor)

def main():
    print 'hi\n'
    X = np.random.randint(10, size=(5,3))
    print X
    Y = np.random.randint(2, size=(5))
    print Y
    dataset = PylearnDataset(X, Y)
    
    print type(dataset)
    print dataset.X
    print dataset.y
    
#     serial.save("dataset.pkl", dataset)

#     t = serial.load("dataset.pkl")
    
main()